//
// Created on 2025/11/5.
//
// Node APIs are not fully supported. To solve the compilation error of the interface cannot be found,
// please include "napi/native_api.h".
#include <fstream>
#include <sstream>
#include <thread>

#include "codec_encoder.h"


static void OnError(OH_AVCodec *codec, int32_t errorCode, void *userData)
{
    // 回调的错误码由开发者判断处理。
    (void)codec;
    (void)errorCode;
    (void)userData;
    CodecUserData *codecData = static_cast<CodecUserData *>(userData);
    codecData->toStop = true;
    PRE_EXPORT_LOGE("code encode error: %{public}d", errorCode);
}

// 解码数据流变化回调OH_AVCodecOnStreamChanged实现。
static void OnStreamChanged(OH_AVCodec *codec, OH_AVFormat *format, void *userData)
{
    PRE_EXPORT_LOGI("OnStreamChanged");
    return;
}

// 解码输入回调OH_AVCodecOnNeedInputBuffer实现。
static void OnNeedInputBuffer(OH_AVCodec *codec, uint32_t index, OH_AVBuffer *buffer, void *userData)
{
    // 输入帧的数据buffer和对应的index送入inQueue队列。
    PRE_EXPORT_LOGI("OnNeedInputBuffer");
    if (userData == nullptr) {
        PRE_EXPORT_LOGE("userData nullptr");
    }
    CodecUserData *codecData = static_cast<CodecUserData *>(userData);
    CacheInfo cacheInfo(index, buffer);
    codecData->inQueue.Enqueue(std::make_shared<CacheInfo>(cacheInfo));
}

// 解码输出回调OH_AVCodecOnNewOutputBuffer实现。
static void OnNewOutputBuffer(OH_AVCodec *codec, uint32_t index, OH_AVBuffer *buffer, void *userData)
{
    PRE_EXPORT_LOGI("OnNewOutputBuffer");
    if (userData == nullptr) {
        PRE_EXPORT_LOGE("userData nullptr");
    }
    CodecUserData *codecData = static_cast<CodecUserData *>(userData);
    // 完成帧的数据buffer和对应的index送入outQueue队列。
    (void)userData;
    CacheInfo cacheInfo(index, buffer);
    codecData->outQueue.Enqueue(std::make_shared<CacheInfo>(cacheInfo));
}

ResultCode CodecEncoder::Configure(uint32_t w, uint32_t h,
                                   OH_AVPixelFormat pixelFormat, uint32_t frameRate)
{
    if (configured_.load()) {
        PRE_EXPORT_LOGI("encoder has been configured");
        return ResultCode::OK;
    }
    configured_.store(true);
    videoEncoder_ = OH_VideoEncoder_CreateByMime(OH_AVCODEC_MIMETYPE_VIDEO_AVC);
    // 配置异步回调，调用OH_VideoDecoder_RegisterCallback接口。
    OH_AVCodecCallback cb = {&OnError, &OnStreamChanged, &OnNeedInputBuffer, &OnNewOutputBuffer};
    // 配置异步回调。
    OH_AVErrCode ret = OH_VideoEncoder_RegisterCallback(videoEncoder_, cb, &codecUserData_);
    if (ret != AV_ERR_OK) {
        PRE_EXPORT_LOGE("OH_VideoDecoder_RegisterCallback failed: %{public}d", ret);
        return ResultCode::ERROR;
    }
    auto format = std::shared_ptr<OH_AVFormat>(OH_AVFormat_Create(), OH_AVFormat_Destroy);
    if (format == nullptr) {
        PRE_EXPORT_LOGE("OH_AVFormat failed: %{public}d", ret);
        return ResultCode::ERROR; 
    }
    // 写入format。
    OH_AVFormat_SetIntValue(format.get(), OH_MD_KEY_WIDTH, w); // 必须配置。
    OH_AVFormat_SetIntValue(format.get(), OH_MD_KEY_HEIGHT, h); // 必须配置。
    OH_AVFormat_SetIntValue(format.get(), OH_MD_KEY_PIXEL_FORMAT, pixelFormat); // 必须配置。
    OH_AVFormat_SetDoubleValue(format.get(), OH_MD_KEY_FRAME_RATE, static_cast<double>(frameRate)); // 必须配置。
    // 配置编码器。
    ret = OH_VideoEncoder_Configure(videoEncoder_, format.get());
    if (ret != AV_ERR_OK) {
        PRE_EXPORT_LOGE("OH_VideoEncoder_Configure failed: %{public}d", ret);
        return ResultCode::ERROR; 
    }
    ret = OH_VideoEncoder_Prepare(videoEncoder_);
    if (ret != AV_ERR_OK) {
        PRE_EXPORT_LOGE("OH_VideoEncoder_Prepare failed: %{public}d", ret);
        return ResultCode::ERROR; 
    }
    return ResultCode::OK;
}

void CodecEncoder::start(std::string cachePath, uint32_t startIndex)
{
    if (!configured_.load() || videoEncoder_ == nullptr) {
        PRE_EXPORT_LOGE("encoder has not been configured or encoder is nullptr");
        return;
    }
    startIndex_ = startIndex;
    cachePath_ = cachePath;
    codecUserData_.toStop = false;
    codecUserData_.inQueue.Flush();
    codecUserData_.outQueue.Flush();
    // 启动编码器，开始编码。
    OH_AVErrCode ret = OH_VideoEncoder_Start(videoEncoder_);
    if (ret != AV_ERR_OK) {
        PRE_EXPORT_LOGE("OH_VideoEncoder_Start failed: %{public}d", ret);
        return;
    }
    return;
}

std::unordered_map<uint64_t, std::string> CodecEncoder::run_until_finish(
    CodecBufferQueue* forwardQueue)
{
    PRE_EXPORT_LOGI("start");
    {
        std::lock_guard<std::mutex> lock(mutex_);
        result_.clear();
    }
    std::thread inThread = std::thread(&CodecEncoder::InputThreadHandler, this, forwardQueue);
    std::thread outThread = std::thread(&CodecEncoder::OutputThreadHandler, this);
    inThread.join();
    outThread.join();
    PRE_EXPORT_LOGI("end");
    return result_;
}

void CodecEncoder::InputThreadHandler(CodecBufferQueue* forwardQueue)
{
    bool isEOS = false;
    uint64_t frameIndex = 0;
    while(!isEOS) {
        if (codecUserData_.toStop) {
            return;
        }
        std::shared_ptr<CacheInfo> inBuffInfo = codecUserData_.inQueue.Dequeue();
        if (inBuffInfo == nullptr || !inBuffInfo->isValid) {
            continue;
        }
        std::shared_ptr<CacheInfo> inDataInfo = forwardQueue->Dequeue();
        if (forwardQueue->empty()) {
            isEOS = true;
        }
        if (inDataInfo == nullptr || !inDataInfo->isValid || !inDataInfo->isReadyToCache()) {
            continue;
        }
        PRE_EXPORT_LOGI("handle %{public}ld", frameIndex);
        frameIndex += 1;
        OH_AVCodecBufferAttr info;
        info.size = inDataInfo->width * inDataInfo->height;
        info.offset = 0;
        info.pts = 0;
        if (isEOS) {
            info.flags = AVCODEC_BUFFER_FLAGS_EOS;
            PRE_EXPORT_LOGI("AVCODEC_BUFFER_FLAGS_EOS");
        }
        uint8_t *src = inDataInfo->inBuff;
        uint8_t *dest = OH_AVBuffer_GetAddr(inBuffInfo->buffer);
        if (src == nullptr || dest == nullptr) {
            PRE_EXPORT_LOGE("get data addr failed");
            continue;
        }
        memcpy(dest, src, info.size);
        auto ret = OH_AVBuffer_SetBufferAttr(inBuffInfo->buffer, &info);
        if (ret != AV_ERR_OK) {
            PRE_EXPORT_LOGE("OH_AVBuffer_SetBufferAttr failed %{public}d", ret);
        }
        ret = OH_VideoEncoder_PushInputBuffer(videoEncoder_, inBuffInfo->index);
        if (ret != AV_ERR_OK) {
            PRE_EXPORT_LOGE("OH_VideoEncoder_PushInputBuffer failed %{public}d", ret);
        }
    } 
}

void CodecEncoder::OutputThreadHandler()
{
    bool isEOS = false;
    uint64_t frameIndex = 0;
    std::unique_ptr<std::ofstream> fd = std::make_unique<std::ofstream>();
    while(!isEOS) {
        if (codecUserData_.toStop) {
            return;
        }
        std::shared_ptr<CacheInfo> outBuffInfo = codecUserData_.outQueue.Dequeue();
        if (outBuffInfo == nullptr || !outBuffInfo->isValid) {
            continue;
        }
        PRE_EXPORT_LOGI("handle %{public}ld", frameIndex);
        OH_AVCodecBufferAttr info;
        auto ret = OH_AVBuffer_GetBufferAttr(outBuffInfo->buffer, &info);
        if (ret != AV_ERR_OK) {
            PRE_EXPORT_LOGE("OH_AVBuffer_GetBufferAttr failed %{public}d", ret);
            continue;
        }
        if (info.flags & AVCODEC_BUFFER_FLAGS_EOS) {
            PRE_EXPORT_LOGI("AVCODEC_BUFFER_FLAGS_EOS");
            isEOS = true;
        }
        uint8_t *addr = OH_AVBuffer_GetAddr(outBuffInfo->buffer);
        if (addr == nullptr) {
            PRE_EXPORT_LOGE("OH_AVBuffer_GetAddr failed %{public}d", ret);
            continue;
        }
        if (fd->is_open()) {
            fd->close();
        }
        {
            std::stringstream ss;
            auto index = frameIndex + startIndex_;
            if (frameIndex == 0) {
                index = 0;
            }
            ss << cachePath_ << "/" << index << ".h264";
            fd->open(ss.str(), std::ios::out | std::ios::binary | std::ios::ate);
            std::lock_guard<std::mutex> lock(mutex_);
            result_[index] = ss.str();
            PRE_EXPORT_LOGI("index %{public}ld - 1 save to %{public}s", index, ss.str().c_str());
        }
        if (frameIndex == 1) {
            void *frameData = nullptr;
            uint32_t frameSize = 0;
            std::lock_guard<std::mutex> lock(mutex_);
            CodecCommon::ReadFile(result_[0], &frameData, &frameSize);
            fd->write(reinterpret_cast<char *>(frameData), frameSize);
        }
        fd->write(reinterpret_cast<char *>(addr), info.size);
        ret = OH_VideoEncoder_FreeOutputBuffer(videoEncoder_, outBuffInfo->index);
        if (ret != AV_ERR_OK) {
            PRE_EXPORT_LOGE("OH_VideoEncoder_FreeOutputBuffer failed %{public}d", ret);
            continue;
        }
        frameIndex += 1;
    }
}

void CodecEncoder::release()
{
    configured_.store(false);
    if (videoEncoder_ != nullptr) {
        OH_VideoEncoder_Stop(videoEncoder_);
        OH_VideoEncoder_Destroy(videoEncoder_);
        videoEncoder_ = nullptr;
    }
    return;
}